package cn.brucemaa.spring_ai_demo.controller;

import cn.brucemaa.spring_ai_demo.tools.DateTimeTools;
import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.transformer.KeywordMetadataEnricher;
import org.springframework.ai.transformer.SummaryMetadataEnricher;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import redis.clients.jedis.JedisPooled;

import java.util.List;
import java.util.Map;

class DemoControllerTest {
    private static final String API_KEY = "sk-eb8cd962b7fe4930ad380d74ff24e6bc";

    public static void main(String[] args) {
        test7();
    }

    private static void test7() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        var chatModel = new DashScopeChatModel(dashScopeApi);
        SummaryMetadataEnricher enricher = new SummaryMetadataEnricher(chatModel,
                List.of(SummaryMetadataEnricher.SummaryType.PREVIOUS, SummaryMetadataEnricher.SummaryType.CURRENT, SummaryMetadataEnricher.SummaryType.NEXT));

        Document doc1 = new Document("Content of document 1");
        Document doc2 = new Document("Content of document 2");

        List<Document> enrichedDocs = enricher.apply(List.of(doc1, doc2));

        // Check the metadata of the enriched documents
        for (Document doc : enrichedDocs) {
            System.out.println("Current summary: " + doc.getMetadata().get("section_summary"));
            System.out.println("Previous summary: " + doc.getMetadata().get("prev_section_summary"));
            System.out.println("Next summary: " + doc.getMetadata().get("next_section_summary"));
        }
    }

    private static void test6() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        var chatModel = new DashScopeChatModel(dashScopeApi);
        KeywordMetadataEnricher enricher = new KeywordMetadataEnricher(chatModel, 5);

        Document doc = new Document("This is a document about artificial intelligence and its applications in modern technology.");

        List<Document> enrichedDocs = enricher.apply(List.of(doc));

        Document enrichedDoc = enrichedDocs.get(0);
        String keywords = (String) enrichedDoc.getMetadata().get("excerpt_keywords");
        System.out.println("Extracted keywords: " + keywords);
    }

    private static void test5() {
        Document doc1 = new Document("This is a long piece of text that needs to be split into smaller chunks for processing.",
                Map.of("source", "example.txt"));
        Document doc2 = new Document("Another document with content that will be split based on token count.",
                Map.of("source", "example2.txt"));

        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> splitDocuments = splitter.apply(List.of(doc1, doc2));

        for (Document doc : splitDocuments) {
            System.out.println("Chunk: " + doc.getFormattedContent());
            System.out.println("Metadata: " + doc.getMetadata());
        }
    }

    private static void test4() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        var chatModel = new DashScopeChatModel(dashScopeApi);
        String response = ChatClient.create(chatModel)
                .prompt("Can you set an alarm 10 minutes from now?")
                .tools(new DateTimeTools())
                .call()
                .content();

        System.out.println(response);
    }

    private static void test3() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        var embeddingModel = new DashScopeEmbeddingModel(dashScopeApi);
        SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
        Document bgDocument = new Document("The World is Big",
                Map.of("country", "Bulgaria"));
        Document nlDocument = new Document("The World is Big",
                Map.of("country", "Netherlands"));

        vectorStore.add(List.of(bgDocument, nlDocument));

        Advisor retrievalAugentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.50)
                        .vectorStore(vectorStore)
                        .build())
                .build();

        var chatModel = new DashScopeChatModel(dashScopeApi);
        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().build()))
                .build();

        // Update filter expression at runtime
        String content = chatClient.prompt()
                .user("这是什么意思？")
                .advisors(retrievalAugentationAdvisor)
                .call()
                .content();
        System.out.println(content);
    }

    private static void test2() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        var dashScopeEmbeddingModel = new DashScopeEmbeddingModel(dashScopeApi);
        try (JedisPooled jedisPooled = new JedisPooled("192.168.5.32", 6379)) {
            RedisVectorStore vectorStore = RedisVectorStore.builder(jedisPooled, dashScopeEmbeddingModel)
                    .indexName("custom-index")                // Optional: defaults to "spring-ai-index"
                    .prefix("custom-prefix")                  // Optional: defaults to "embedding:"
                    .metadataFields(                         // Optional: define metadata fields for filtering
                            RedisVectorStore.MetadataField.tag("country"),
                            RedisVectorStore.MetadataField.numeric("year"))
                    .initializeSchema(true)                   // Optional: defaults to false
                    .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy
                    .build();

            Document bgDocument = new Document("The World is Big",
                    Map.of("country", "Bulgaria"));
            Document nlDocument = new Document("The World is Big",
                    Map.of("country", "Netherlands"));

            vectorStore.add(List.of(bgDocument, nlDocument));

            Filter.Expression filterExpression = new Filter.Expression(
                    Filter.ExpressionType.EQ,
                    new Filter.Key("country"),
                    new Filter.Value("Bulgaria")
            );

            vectorStore.delete(filterExpression);
            SearchRequest request = SearchRequest.builder()
                    .query("World")
                    .filterExpression("country == 'Bulgaria'")
                    .build();
            List<Document> results = vectorStore.similaritySearch(request);
            System.out.println(results);

            request = SearchRequest.builder()
                    .query("World")
                    .filterExpression("country == 'Netherlands'")
                    .build();
            results = vectorStore.similaritySearch(request);
            System.out.println(results);
        }
    }

    private static void test1() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        var dashScopeEmbeddingModel = new DashScopeEmbeddingModel(dashScopeApi);
        SimpleVectorStore vectorStore = SimpleVectorStore.builder(dashScopeEmbeddingModel).build();

        Document bgDocument = new Document("The World is Big",
                Map.of("country", "Bulgaria"));
        Document nlDocument = new Document("The World is Big",
                Map.of("country", "Netherlands"));

        vectorStore.add(List.of(bgDocument, nlDocument));

        Filter.Expression filterExpression = new Filter.Expression(
                Filter.ExpressionType.EQ,
                new Filter.Key("country"),
                new Filter.Value("Bulgaria")
        );

        vectorStore.delete(List.of(bgDocument.getId()));

        SearchRequest request = SearchRequest.builder()
                .query("World")
                .filterExpression("country == 'Bulgaria'")
                .build();
        List<Document> results = vectorStore.similaritySearch(request);
        System.out.println(results);

        request = SearchRequest.builder()
                .query("World")
                .filterExpression("country == 'Netherlands'")
                .build();
        results = vectorStore.similaritySearch(request);
        System.out.println(results);
    }
}